#!/usr/bin/env python

# Author: Edward Roualdes
#         [2012.11.20]
#         University of Kentucky

from Bio import SeqIO
import getopt
import sys
import os
import errno
import re
import shutil


def parseCommandLine(argv):
    """
    extract command line options provided.
    offer help / suggestions if arguments given are incorrect.
    """
    
    outputfile = ""
    try:
        options, remainder = getopt.getopt(sys.argv[1:], 
                                           'hc:v', 
                                           ['help', 'configfile=', 'version'])
    except getopt.GetoptError as err:
        print str(err)
        print "Try instead '$ python NextAllele.py -c path/to/CONFIGFILE'"
        sys.exit(2)
    for opt, arg in options:
        if opt in ('-h', '--help'):
            print "Usage"
            print "\tpython NextAllele.py [FLAGS] [CONFIGFILE]\n"
            print "Flags"
            print "{:<20} {:<}".format("\t-h, --help", "print this help message and exit")
            print "{:<20} {:<}".format("\t-v, --version", "print version number and exit")
            print "{:<20} {:<}".format("\t-c, --configfile", "specify CONFIGFILE: -c path/to/CONFIGFILE")
            sys.exit()
        elif opt in ('-v', '--version'):
            print "NextAllele v0.1"
            sys.exit()
        elif opt in ('-c', '--configfile'):
            outputfile = arg
    print "\nCONFIGFILE =", os.path.abspath(outputfile)
    return outputfile


def getNames(FILE):
    """get read names from FILE.
    return set of names.
    """

    Names = set()
    file = open(FILE, "rU")
    for seq in SeqIO.parse(file, "fasta"): 
        Names.add(seq.name)
    file.close()

    return Names


def mkdir(path, overwrite, opt = None):
    """
    make directory with the option to overwrite (=True) existing directory tree, or exit program to specify a new directory. If overwrite == True, specify an opt (= some string)
    """

    try:
        os.makedirs(path)       # just try to make it
    except OSError as exc:      # catch exception
        if exc.errno == errno.EEXIST:
            if overwrite:
                # if exists, choice
                inp = raw_input(path + " already exists, do you want to overwrite it and all of its subdirectories? ")
                res = re.match("[Yy][Ee][Ss]|[Yy]", inp)
                if res:             # delete directory and re-make it
                    shutil.rmtree(path)
                    os.makedirs(path)
                else: # give option kill NextAllele and specify a new OUTPUTDIRECTORY
                    print opt + "and rerun NextAllele."
                    sys.exit()
            else:
                pass
        else:
            raise
    return


def writeSeqs(sortedBarcodes, outDir):
    """write sorted sequences to files: 
    one folder per individual, one file per locus
    """

    for inds, reads in sortedBarcodes.iteritems():
        for read in reads:      # for each read
            individual = inds
            locus = read.annotations["Locus"]

            # make dir each individual w/o overwriting
            mkdir(os.path.join(outDir, individual), False)

            # one file per locus
            seqFile = open(os.path.join(outDir, individual, locus+'.fna'), 'a')

            # write sequences
            SeqIO.write(read, seqFile, "fasta")
            seqFile.close()
    print "Wrote alignments to file..."


# save fasta files in directory to dictionaries, yield one at a time
def files2dict(DIRECTORY, grep, pat = None):
    """
    walks directory and iteratively yields a dictionary (k = sequence name, v = sequence) for each grep (a str pattern) matched fasta file within DIRECTORY and any subdirectories, indiv, locus names found via directory's path structure: directory => (indiv1 => loci1), (indiv1 => loci2), (indiv1 => lociK) ... , (indivN => lociM)
    """
    grep = re.compile(grep)
    if not pat:
        # use 'modified' if not supplied alternative key word
        pat = 'modified'

    for root, dirs, files in os.walk(DIRECTORY):
        # find and use only files matching grep
        for fname in filter(grep.match, files):

            f = os.path.join(root, fname)  # use absolute path names
            indiv = os.path.basename(root) # save indiv name
            locus = os.path.splitext(os.path.splitext(fname)[0])[0] # save locus name
            data = {}
            # look for "modified" version of each aligned locus file 
            for r, ds, fs in os.walk(root):
                for fsname in filter(lambda x: locus in x, fs):
                    if pat in fsname: # if it exists
                        f = os.path.join(root, fsname) # read from it

            alignedFile = open(f, "rU")
            data = SeqIO.to_dict(SeqIO.parse(alignedFile, "fasta"))
            alignedFile.close()
            yield data, indiv, locus

if __name__ == "__main__":
    parseCommandLine(sys.argv[1:])
